package org.lop.modules.utils;

import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.sql.SQLException;

import javax.sql.DataSource;

import org.apache.commons.lang3.ArrayUtils;
import org.apache.commons.lang3.StringUtils;
import org.apache.commons.lang3.Validate;
import org.dbunit.DatabaseUnitException;
import org.dbunit.database.DatabaseConnection;
import org.dbunit.database.IDatabaseConnection;
import org.dbunit.dataset.IDataSet;
import org.dbunit.dataset.xml.FlatXmlDataSet;
import org.dbunit.dataset.xml.FlatXmlDataSetBuilder;
import org.dbunit.operation.DatabaseOperation;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.core.io.DefaultResourceLoader;
import org.springframework.core.io.ResourceLoader;

/**
 * DBUnit工具.
 * 
 * @author 潘瑞峥
 * @date 2013-6-7
 */
public class DbUnitUtils {

	private static Logger logger = LoggerFactory.getLogger( DbUnitUtils.class );

	private static ResourceLoader resourceLoader = new DefaultResourceLoader();

	/**
	 * 将数据库数据写入XML文件.
	 */
	public static void writeData( DataSource dataSource, String xmlFilePath, String... tableNames ) throws Exception {
		Validate.notBlank( xmlFilePath, "xmlFilePath不能为空" );

		String suffix = ".xml";
		IDataSet dataSet = null;
		IDatabaseConnection connection = new DatabaseConnection( dataSource.getConnection() );
		/* 导出全部. */
		if ( ArrayUtils.isEmpty( tableNames ) ) {
			dataSet = connection.createDataSet();
			FlatXmlDataSet.write( dataSet, createOutputStream( xmlFilePath, suffix, null ) );
		} else {
			for ( String tableName : tableNames ) {
				dataSet = connection.createDataSet( new String[] { tableName } );
				FlatXmlDataSet.write( dataSet, createOutputStream( xmlFilePath, suffix, tableName ) );
			}
		}
	}

	/**
	 * 根据参数创建OutputStream.
	 */
	private static FileOutputStream createOutputStream( String path, String suffix, String name ) throws Exception {
		String result = null;
		if ( StringUtils.endsWithIgnoreCase( path, suffix ) ) {
			result = StringUtils.substring( path, 0, StringUtils.indexOfIgnoreCase( path, suffix ) );
		} else {
			result = path;
		}

		if ( StringUtils.isNotBlank( name ) ) {
			result += "_" + name;
		} else {
			result += "_all";
		}

		result += suffix;

		return new FileOutputStream( new File( result ) );
	}

	/**
	 * 清除并插入XML数据文件到数据库.
	 * 
	 * XML数据文件中涉及的表在插入数据前会先进行清除.
	 * 
	 * @param xmlFilePaths
	 *            符合Spring Resource路径格式的文件列表.
	 */
	public static void loadData( DataSource dataSource, String... xmlFilePaths ) throws Exception {
		execute( DatabaseOperation.CLEAN_INSERT, dataSource, xmlFilePaths );
	}

	/**
	 * 插入XML数据文件到数据库.
	 */
	public static void appendData( DataSource dataSource, String... xmlFilePaths ) throws Exception {
		execute( DatabaseOperation.INSERT, dataSource, xmlFilePaths );
	}

	/**
	 * 在数据库中删除XML数据文件中涉及的表的数据.
	 */
	public static void removeData( DataSource dataSource, String... xmlFilePaths ) throws Exception {
		execute( DatabaseOperation.DELETE_ALL, dataSource, xmlFilePaths );
	}

	/**
	 * 按DBUnit Operation执行XML数据文件的数据.
	 * 
	 * @param xmlFilePaths
	 *            符合Spring Resource路径格式的文件列表.
	 */
	private static void execute( DatabaseOperation operation, DataSource dataSource, String... xmlFilePaths ) throws DatabaseUnitException,
			SQLException {
		Validate.notNull( xmlFilePaths, "xmlFilePaths不能为空" );

		IDatabaseConnection connection = new DatabaseConnection( dataSource.getConnection(), "" );
		for ( String xmlPath : xmlFilePaths ) {
			try {
				InputStream input = resourceLoader.getResource( xmlPath ).getInputStream();
				IDataSet dataSet = new FlatXmlDataSetBuilder().setColumnSensing( true ).build( input );
				operation.execute( connection, dataSet );
			} catch ( IOException e ) {
				logger.warn( xmlPath + " file not found", e );
			}
		}
	}

}